import os
import pandas as pd
import cv2
from pathlib import Path
import glob
import sys
sys.path.append('..')
from utils.json import json_to_image


def get_crop_locs(edge_len, base_size):
    edge_times =  edge_len // base_size
    edge_crops = []
    edge_locs = []

    if edge_times == 0:
        edge_times == 1
        edge_crops = [edge_len]
        edge_locs = [0]
    else:
        carry = edge_len % base_size
        
        if carry >= 100:
            edge_locs = [i*base_size for i in range(edge_times)]
            edge_locs.append(edge_len-base_size)
            edge_crops = [base_size] * (edge_times + 1)
        else:
            base_size_dilate =  base_size + carry // edge_times
            edge_locs = [i*base_size_dilate for i in range(edge_times-1)]
            last_base_size = base_size_dilate + carry % edge_times
            edge_locs.append((edge_times-1) * base_size_dilate)
            edge_crops = [base_size_dilate] * (edge_times -1) + [last_base_size]
    return edge_locs, edge_crops

def adaptive_crop(img, base_size=1024):
    h, w = img.shape[:2]

    ys, h_crops = get_crop_locs(h, base_size)
    xs, w_crops = get_crop_locs(w, base_size)
    return xs, ys, h_crops, w_crops

def mid_crop(img):
    h, w = img.shape[:2]
    xs, ys = [0, w//2], [0, h//2]
    w_crops = [w//2, w-w//2]
    h_crops = [h//2, h-h//2]

    return xs, ys, h_crops, w_crops

    

def img_split(img_path, mask_root, img_dst_root, mask_dst_root, mask_suffix, no_defect_code):
    try:
        img = cv2.imread(img_path, 1)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        code = Path(img_path).parent.name
        name_stem = Path(img_path).stem
        mask_path = os.path.join(mask_root, code, name_stem+mask_suffix)
        if (code!=no_defect_code) and (not os.path.exists(mask_path)):
            return

        if mask_suffix == '.png':
            mask = cv2.imread(mask_path, 0)
        elif mask_suffix == '.json':
            mask = json_to_image(img_shape=img.shape, json_path=mask_path)
            mask *= 255

        # xs, ys, h_crops, w_crops = adaptive_crop(img)
        xs, ys, h_crops, w_crops = mid_crop(img)

        index = 1
        for x, w_crop in zip(xs, w_crops):
            for y, h_crop in zip(ys, h_crops):
                cur_img = img[y:y+h_crop, x:x+w_crop]
                cur_mask = mask[y:y+h_crop, x:x+w_crop]

                if 255 in cur_mask:
                    cur_code = code
                    os.makedirs(os.path.join(img_dst_root, cur_code), exist_ok=True)
                    cv2.imwrite(os.path.join(img_dst_root, cur_code, name_stem+"_"+str(index)+'.jpg'), cur_img)
                    os.makedirs(os.path.join(mask_dst_root, cur_code), exist_ok=True)
                    cv2.imwrite(os.path.join(mask_dst_root, cur_code, name_stem+"_"+str(index)+'.png'), cur_mask)
                else:
                    cur_code = no_defect_code
                    os.makedirs(os.path.join(img_dst_root, cur_code), exist_ok=True)
                    cv2.imwrite(os.path.join(img_dst_root, cur_code, name_stem+"_"+str(index)+'.jpg'), cur_img)
                
                index += 1
    except:
        return

def main():
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    # img_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409'
    # mask_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409_mask'
    
    # img_dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_img_0409'
    # mask_dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_mask_0409'

    img_ps = [r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6002/all_0401']

    mask_suffix = '.json'
    no_defect_code = 'TGXID'

    for img_root in img_ps:
        mask_root = img_root[:]
        img_dst_root = img_root + "_crop"
        mask_dst_root = img_dst_root

        df = pd.DataFrame()
        df['img_path'] = glob.glob(os.path.join(img_root, "*/*.jpg"))
        print()
        print(len(df))

        if is_pandarallel:
            df.parallel_apply(lambda x:img_split(img_path=x['img_path'], mask_root=mask_root, img_dst_root=img_dst_root, mask_dst_root=mask_dst_root, mask_suffix=mask_suffix, no_defect_code=no_defect_code), axis=1)
        else:
            df.apply(lambda x:img_split(img_path=x['img_path'], mask_root=mask_root, img_dst_root=img_dst_root, mask_dst_root=mask_dst_root, mask_suffix=mask_suffix, no_defect_code=no_defect_code), axis=1)



if __name__=='__main__':
    main()
